import json
import os
import argparse
# from collections import defaultdict
from tqdm import tqdm
from PIL import Image  # 用于动态获取图像宽度和高度

def yolo_to_coco(input_dir, output_dir):
    """
    将YOLO格式的数据集转换为COCO格式的数据集
    :param input_dir: 包含YOLO格式标注文件和图像文件的目录
    :param output_dir: 输出目录，用于保存COCO格式的JSON文件
    """
    # 初始化COCO格式的JSON结构
    coco_output = {
        "images": [],
        "annotations": [],
        "categories": []
    }

    # 读取类别名称文件（假设类别文件名为classes.txt）
    classes_path = os.path.join(input_dir, "classes.txt")
    if not os.path.exists(classes_path):
        raise FileNotFoundError(f"类别文件 {classes_path} 不存在！")

    with open(classes_path, "r") as f:
        classes = [line.strip() for line in f.readlines()]

    # 添加类别信息到COCO格式
    for idx, class_name in enumerate(classes):
        category = {
            "id": idx + 1,  # COCO类别ID从1开始
            "name": class_name,
            "supercategory": "object"
        }
        coco_output["categories"].append(category)

    # 遍历YOLO标注文件
    annotation_files = [f for f in os.listdir(input_dir) if f.endswith(".txt") and f != "classes.txt"]
    image_files = [f.replace(".txt", ".png") for f in annotation_files]  # 假设图像格式为jpg
    # print(annotation_files)
    for annotation_file, image_file in tqdm( zip(annotation_files, image_files) ):
        annotation_path = os.path.join(input_dir, annotation_file)
        image_path = os.path.join(input_dir, image_file)

        # 检查图像文件是否存在
        if not os.path.exists(image_path):
            print(f"图像文件 {image_path} 不存在，跳过...")
            continue

        # 动态获取图像宽度和高度
        with Image.open(image_path) as img:
            image_width, image_height = img.size

        # 添加图像信息到COCO格式
        image_id = len(coco_output["images"]) + 1
        image_info = {
            "id": image_id,
            "file_name": image_file,
            "width": image_width,
            "height": image_height
        }
        coco_output["images"].append(image_info)

        # 读取YOLO标注文件
        with open(annotation_path, "r") as f:
            for line in f.readlines():
                parts = line.strip().split()
                if len(parts) != 5:
                    print(f"标注文件 {annotation_path} 格式错误，跳过...")
                    continue

                # 解析YOLO格式的标注
                class_id, x_center, y_center, width, height = map(float, parts)
                class_id = int(class_id)

                # 转换为COCO格式的标注
                x_min = (x_center - width / 2.0) * image_width
                y_min = (y_center - height / 2.0) * image_height
                bbox_width = width * image_width
                bbox_height = height * image_height

                annotation = {
                    "id": len(coco_output["annotations"]) + 1,
                    "image_id": image_id,
                    "category_id": class_id + 1,  # COCO类别ID从1开始
                    "bbox": [x_min, y_min, bbox_width, bbox_height],
                    "area": bbox_width * bbox_height,
                    "segmentation": [],
                    "iscrowd": 0
                }
                coco_output["annotations"].append(annotation)

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 写入COCO格式的JSON文件
    output_json_path = os.path.join(output_dir, "dataset.json")
    with open(output_json_path, "w") as f:
        json.dump(coco_output, f,ensure_ascii=False, indent=4)

    print(f"转换完成！COCO格式的标注文件已保存到 {output_json_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="将YOLO格式数据集转换为COCO格式数据集")
    parser.add_argument("--input_dir", type=str, required=False, help="包含YOLO格式标注文件和图像文件的目录",default="./test")
    parser.add_argument("--output_dir", type=str, required=False, help="输出目录，用于保存COCO格式的JSON文件",default="./coco_data")
    # 输出帮助信息
    parser.print_help()
    
    args = parser.parse_args([])
    yolo_to_coco(args.input_dir, args.output_dir)


